{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-11-03T03:33:24.108347Z",
     "start_time": "2025-11-03T03:33:23.011289Z"
    }
   },
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "import math"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-03T03:33:29.794301Z",
     "start_time": "2025-11-03T03:33:29.783379Z"
    }
   },
   "cell_type": "code",
   "source": [
    "random_torch=torch.rand(4,4)\n",
    "print(random_torch)"
   ],
   "id": "2645f4345627449c",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.4693, 0.8979, 0.7669, 0.0991],\n",
      "        [0.9914, 0.4338, 0.0196, 0.3992],\n",
      "        [0.7096, 0.8356, 0.1954, 0.3299],\n",
      "        [0.5330, 0.5268, 0.0604, 0.9321]])\n"
     ]
    }
   ],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-03T03:33:33.764550Z",
     "start_time": "2025-11-03T03:33:33.760085Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from torch import tensor\n",
    "# 将输入的词汇表索引转换为指定维度的 Embedding\n",
    "\n",
    "class TokenEmbedding(nn.Embedding):\n",
    "    def __init__(self, vocab_size, d_model):\n",
    "        super(TokenEmbedding, self).__init__(vocab_size, d_model, padding_idx=1)"
   ],
   "id": "1775f4b2cf18e5d5",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-03T03:33:36.520742Z",
     "start_time": "2025-11-03T03:33:36.512021Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class PositionalEmbedding(nn.Module):\n",
    "    def __init__(self, d_model, max_len=5000, device='cpu'):\n",
    "        super(PositionalEmbedding, self).__init__()\n",
    "        self.encoding = torch.zeros(max_len, d_model,device=device)\n",
    "        self.encoding.requires_grad = False\n",
    "        pos = torch.arange(0, max_len, device=device)\n",
    "        pos = pos.float().unsqueeze(dim=1)\n",
    "        _2i = torch.arange(0, d_model, step=2, device=device).float()\n",
    "        self.encoding[:, 0::2] = torch.sin(pos / (10000 ** (_2i / d_model)))\n",
    "        self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (_2i / d_model)))\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size, seq_len = x.size()\n",
    "        return self.encoding[:seq_len, :]"
   ],
   "id": "d2e1ed20dd4beae5",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-11-03T03:34:14.627432Z",
     "start_time": "2025-11-03T03:34:14.612119Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class TransformerEmbedding(nn.Module):\n",
    "    def __init__(self, vocab_size, d_model, max_len=5000, drop_prob=0.1, device='cpu'):\n",
    "        super(TransformerEmbedding, self).__init__()\n",
    "        self.token_embedding = TokenEmbedding(vocab_size, d_model)\n",
    "        self.position_embedding = PositionalEmbedding(d_model, max_len, device)\n",
    "        self.embedding_dropout = nn.Dropout(p = drop_prob)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.token_embedding(x) + self.position_embedding(x)\n",
    "        out = self.embedding_dropout(out)\n",
    "        return out"
   ],
   "id": "59207f39a4fa0315",
   "outputs": [],
   "execution_count": 7
  },
  {
   "metadata": {
    "jupyter": {
     "is_executing": true
    },
    "ExecuteTime": {
     "start_time": "2025-11-03T03:38:13.281658Z"
    }
   },
   "cell_type": "code",
   "source": [
    "vocab_size = 10000  # 词汇表大小\n",
    "d_model = 512       # 模型维度\n",
    "max_len = 512       # 最大序列长度\n",
    "batch_size = 2      # 批次大小\n",
    "seq_len = 10        # 序列长度\n",
    "\n",
    "# 创建 TransformerEmbedding 实例\n",
    "embedding = TransformerEmbedding(vocab_size, d_model, max_len, device='cpu')\n",
    "\n",
    "# 创建模拟的输入数据 (批次大小, 序列长度)\n",
    "# 这里是词汇索引，值应在 [0, vocab_size) 范围内\n",
    "input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))\n",
    "\n",
    "# 获取嵌入表示\n",
    "embedded_output = embedding(input_ids)\n",
    "\n",
    "print(f\"Input shape: {input_ids.shape}\")           # 应该是 [2, 10]\n",
    "print(f\"Output shape: {embedded_output.shape}\")    # 应该是 [2, 10, 512]\n",
    "print(f\"Output sample:\\n{embedded_output[0, :2, :5]}\")  # 显示部分输出示例"
   ],
   "id": "b1be7927ce097f4",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
